from SpongeExt import enhanced_sponge

import pysages
from pysages.colvars import DihedralAngle
from SpongeExt.pysages import SinkMetaDynamics
from SpongeExt.pysages import MetaDynamics

# def phi_psi():
#     from numpy import pi
#     cvs = [DihedralAngle([4, 6, 8, 14]), DihedralAngle([6, 8, 14, 16])]
#     grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(50, 50), periodic=True)
#     method = SinkMetaDynamics(cv=cvs,
#                               grid=grid,
#                               E_depth=25.0,   # 下沉深度(kJ/mol)
#                               sigma=0.2,       # 高斯宽度
#                               gamma=10.0,      # Well-Tempered因子
#                               cv_bounds=[(-1.0, 1.0), (-1.0, 1.0)]  # 限制phi/psi在[-1,1]弧度内
#                               )
#     return method

# def phi_psi():
#     from numpy import pi
#     cvs = [DihedralAngle([4, 6, 8, 14])]
#     grid = pysages.Grid(lower=(-pi,), upper=(pi,), shape=(50,), periodic=True)
#     method = SinkMetaDynamics(cv=cvs,
#                               grid=grid,
#                               E_depth=25.0,   # 下沉深度(kJ/mol)
#                               sigma=0.2,       # 高斯宽度
#                               gamma=10.0,      # Well-Tempered因子
#                               cv_bounds=[(-3.0, -2.7)],  # 限制phi/psi在[-1,1]弧度内
#                               stride=50
#                               )
#     return method

def phi_psi():
    from numpy import pi
    cvs = [DihedralAngle([4, 6, 8, 14])]
    grid = pysages.Grid(lower=(-pi,), upper=(pi,), shape=(50,), periodic=True)
    method = MetaDynamics(cv=cvs,
                              grid=grid,
                              sigma=0.5,       # 高斯宽度
                              gamma=10.0,      # Well-Tempered因子
                              cv_bounds=[(-pi, pi)],  # 限制phi/psi在[-1,1]弧度内
                              stride=20
                              )
    return method

pysages_method = phi_psi()
Calculate_Force, Mdout_Print = enhanced_sponge(pysages_method)
